#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
任务分配和负载均衡模块
实现基础的任务分配和负载均衡功能
"""

import time
import threading
import logging
from typing import Dict, List, Optional, Any
from dataclasses import dataclass
from enum import Enum
import json

logger = logging.getLogger(__name__)


class TaskStatus(Enum):
    """任务状态枚举"""
    QUEUED = "queued"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"
    TIMEOUT = "timeout"


class ServerStatus(Enum):
    """服务器状态枚举"""
    HEALTHY = "healthy"
    UNHEALTHY = "unhealthy"
    OVERLOADED = "overloaded"
    MAINTENANCE = "maintenance"


@dataclass
class Task:
    """任务数据类"""
    task_id: str
    user_id: str
    client_ip: str
    task_data: Dict[str, Any]
    submit_time: float
    start_time: Optional[float] = None
    complete_time: Optional[float] = None
    status: TaskStatus = TaskStatus.QUEUED
    assigned_server: Optional[str] = None
    error_message: Optional[str] = None
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            'task_id': self.task_id,
            'user_id': self.user_id,
            'client_ip': self.client_ip,
            'task_data': self.task_data,
            'submit_time': self.submit_time,
            'start_time': self.start_time,
            'complete_time': self.complete_time,
            'status': self.status.value,
            'assigned_server': self.assigned_server,
            'error_message': self.error_message
        }


@dataclass
class Server:
    """服务器数据类"""
    server_id: str
    region: str
    gpu_type: str
    status: ServerStatus = ServerStatus.HEALTHY
    cpu_usage: float = 0.0
    gpu_usage: float = 0.0
    memory_usage: float = 0.0
    current_tasks: int = 0
    max_concurrent_tasks: int = 10
    last_health_check: float = 0.0
    created_time: float = 0.0
    endpoint: str = ""
    
    def is_available(self) -> bool:
        """检查服务器是否可用"""
        return (self.status == ServerStatus.HEALTHY and 
                self.current_tasks < self.max_concurrent_tasks)
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            'server_id': self.server_id,
            'region': self.region,
            'gpu_type': self.gpu_type,
            'status': self.status.value,
            'cpu_usage': self.cpu_usage,
            'gpu_usage': self.gpu_usage,
            'memory_usage': self.memory_usage,
            'current_tasks': self.current_tasks,
            'max_concurrent_tasks': self.max_concurrent_tasks,
            'last_health_check': self.last_health_check,
            'created_time': self.created_time,
            'endpoint': self.endpoint
        }


class LoadBalancer:
    """负载均衡器"""
    
    def __init__(self):
        self.servers: Dict[str, Server] = {}
        self.tasks: Dict[str, Task] = {}
        self.task_lock = threading.Lock()
        self.server_lock = threading.Lock()
        
    def add_server(self, server: Server) -> None:
        """添加服务器"""
        with self.server_lock:
            self.servers[server.server_id] = server
            logger.info(f"添加服务器: {server.server_id} ({server.region}, {server.gpu_type})")
    
    def remove_server(self, server_id: str) -> bool:
        """移除服务器"""
        with self.server_lock:
            if server_id in self.servers:
                server = self.servers[server_id]
                # 重新分配该服务器的任务
                self._redistribute_server_tasks(server_id)
                del self.servers[server_id]
                logger.info(f"移除服务器: {server_id}")
                return True
            return False
    
    def update_server_status(self, server_id: str, status: ServerStatus, 
                           cpu_usage: float = None, gpu_usage: float = None,
                           memory_usage: float = None, current_tasks: int = None) -> bool:
        """更新服务器状态"""
        with self.server_lock:
            if server_id in self.servers:
                server = self.servers[server_id]
                server.status = status
                server.last_health_check = time.time()
                
                if cpu_usage is not None:
                    server.cpu_usage = cpu_usage
                if gpu_usage is not None:
                    server.gpu_usage = gpu_usage
                if memory_usage is not None:
                    server.memory_usage = memory_usage
                if current_tasks is not None:
                    server.current_tasks = current_tasks
                
                logger.debug(f"更新服务器状态: {server_id} -> {status.value}")
                return True
            return False
    
    def select_best_server(self, preferred_regions: List[str] = None) -> Optional[Server]:
        """
        选择最佳服务器
        
        Args:
            preferred_regions: 优先选择的地域列表
            
        Returns:
            最佳服务器实例，如果没有可用服务器则返回None
        """
        with self.server_lock:
            available_servers = [s for s in self.servers.values() if s.is_available()]
            
            if not available_servers:
                return None
            
            # 如果有指定优先地域，优先选择这些地域的服务器
            if preferred_regions:
                preferred_servers = [s for s in available_servers if s.region in preferred_regions]
                if preferred_servers:
                    available_servers = preferred_servers
            
            # 按负载和GPU类型选择最佳服务器
            # 优先选择负载最低且GPU类型最便宜的服务器
            best_server = min(available_servers, key=lambda s: (
                s.current_tasks / s.max_concurrent_tasks,  # 负载比例
                s.gpu_usage,  # GPU使用率
                s.cpu_usage,  # CPU使用率
                s.memory_usage  # 内存使用率
            ))
            
            return best_server
    
    def assign_task(self, task: Task) -> bool:
        """
        分配任务到服务器
        
        Args:
            task: 任务实例
            
        Returns:
            是否成功分配
        """
        server = self.select_best_server()
        
        if server is None:
            logger.warning(f"没有可用服务器分配任务: {task.task_id}")
            return False
        
        with self.task_lock:
            task.assigned_server = server.server_id
            task.status = TaskStatus.PROCESSING
            task.start_time = time.time()
            self.tasks[task.task_id] = task
        
        with self.server_lock:
            server.current_tasks += 1
        
        logger.info(f"任务已分配: {task.task_id} -> {server.server_id} ({server.region})")
        return True
    
    def complete_task(self, task_id: str, success: bool = True, 
                     error_message: str = None) -> bool:
        """
        完成任务
        
        Args:
            task_id: 任务ID
            success: 是否成功
            error_message: 错误信息
            
        Returns:
            是否成功完成
        """
        with self.task_lock:
            if task_id not in self.tasks:
                logger.warning(f"任务不存在: {task_id}")
                return False
            
            task = self.tasks[task_id]
            task.complete_time = time.time()
            task.status = TaskStatus.COMPLETED if success else TaskStatus.FAILED
            task.error_message = error_message
            
            # 释放服务器资源
            if task.assigned_server:
                with self.server_lock:
                    if task.assigned_server in self.servers:
                        self.servers[task.assigned_server].current_tasks -= 1
            
            logger.info(f"任务完成: {task_id} -> {task.status.value}")
            return True
    
    def get_server_metrics(self, server_id: str) -> Optional[Dict[str, Any]]:
        """获取服务器指标"""
        with self.server_lock:
            if server_id in self.servers:
                server = self.servers[server_id]
                return {
                    'server_id': server_id,
                    'status': server.status.value,
                    'cpu_usage': server.cpu_usage,
                    'gpu_usage': server.gpu_usage,
                    'memory_usage': server.memory_usage,
                    'current_tasks': server.current_tasks,
                    'max_concurrent_tasks': server.max_concurrent_tasks,
                    'load_ratio': server.current_tasks / server.max_concurrent_tasks,
                    'last_health_check': server.last_health_check,
                    'uptime': time.time() - server.created_time
                }
            return None
    
    def get_system_metrics(self) -> Dict[str, Any]:
        """获取系统整体指标"""
        with self.server_lock:
            total_servers = len(self.servers)
            healthy_servers = len([s for s in self.servers.values() if s.status == ServerStatus.HEALTHY])
            available_servers = len([s for s in self.servers.values() if s.is_available()])
            total_capacity = sum(s.max_concurrent_tasks for s in self.servers.values())
            used_capacity = sum(s.current_tasks for s in self.servers.values())
            
            # 按地域统计
            region_stats = {}
            for server in self.servers.values():
                if server.region not in region_stats:
                    region_stats[server.region] = {
                        'total_servers': 0,
                        'healthy_servers': 0,
                        'available_servers': 0,
                        'total_capacity': 0,
                        'used_capacity': 0
                    }
                
                stats = region_stats[server.region]
                stats['total_servers'] += 1
                if server.status == ServerStatus.HEALTHY:
                    stats['healthy_servers'] += 1
                if server.is_available():
                    stats['available_servers'] += 1
                stats['total_capacity'] += server.max_concurrent_tasks
                stats['used_capacity'] += server.current_tasks
        
        with self.task_lock:
            total_tasks = len(self.tasks)
            processing_tasks = len([t for t in self.tasks.values() if t.status == TaskStatus.PROCESSING])
            queued_tasks = len([t for t in self.tasks.values() if t.status == TaskStatus.QUEUED])
            completed_tasks = len([t for t in self.tasks.values() if t.status == TaskStatus.COMPLETED])
            failed_tasks = len([t for t in self.tasks.values() if t.status == TaskStatus.FAILED])
        
        return {
            'servers': {
                'total': total_servers,
                'healthy': healthy_servers,
                'available': available_servers,
                'total_capacity': total_capacity,
                'used_capacity': used_capacity,
                'utilization_rate': used_capacity / total_capacity if total_capacity > 0 else 0
            },
            'tasks': {
                'total': total_tasks,
                'processing': processing_tasks,
                'queued': queued_tasks,
                'completed': completed_tasks,
                'failed': failed_tasks
            },
            'regions': region_stats
        }
    
    def _redistribute_server_tasks(self, server_id: str) -> int:
        """重新分配服务器任务"""
        redistributed_count = 0
        
        with self.task_lock:
            tasks_to_redistribute = [
                task for task in self.tasks.values() 
                if task.assigned_server == server_id and task.status == TaskStatus.PROCESSING
            ]
            
            for task in tasks_to_redistribute:
                task.status = TaskStatus.QUEUED
                task.assigned_server = None
                task.start_time = None
                redistributed_count += 1
        
        logger.info(f"重新分配任务: {server_id} -> {redistributed_count}个任务")
        return redistributed_count
    
    def cleanup_completed_tasks(self, max_age_hours: int = 24) -> int:
        """清理已完成的任务"""
        current_time = time.time()
        max_age_seconds = max_age_hours * 3600
        cleaned_count = 0
        
        with self.task_lock:
            tasks_to_remove = []
            
            for task_id, task in self.tasks.items():
                if (task.status in [TaskStatus.COMPLETED, TaskStatus.FAILED] and
                    current_time - task.complete_time > max_age_seconds):
                    tasks_to_remove.append(task_id)
            
            for task_id in tasks_to_remove:
                del self.tasks[task_id]
                cleaned_count += 1
        
        if cleaned_count > 0:
            logger.info(f"清理已完成任务: {cleaned_count}个")
        
        return cleaned_count


class TaskDistributor:
    """任务分发器"""
    
    def __init__(self, load_balancer: LoadBalancer):
        self.load_balancer = load_balancer
        self.running = False
        self.worker_thread = None
        
    def start(self) -> None:
        """启动任务分发器"""
        if self.running:
            return
        
        self.running = True
        self.worker_thread = threading.Thread(target=self._worker_loop, daemon=True)
        self.worker_thread.start()
        logger.info("任务分发器已启动")
    
    def stop(self) -> None:
        """停止任务分发器"""
        self.running = False
        if self.worker_thread:
            self.worker_thread.join(timeout=5)
        logger.info("任务分发器已停止")
    
    def _worker_loop(self) -> None:
        """工作循环"""
        while self.running:
            try:
                # 这里应该从公平队列获取任务
                # 暂时使用模拟逻辑
                self._process_pending_tasks()
                time.sleep(1)  # 每秒检查一次
            except Exception as e:
                logger.error(f"任务分发器错误: {e}")
                time.sleep(5)
    
    def _process_pending_tasks(self) -> None:
        """处理待处理任务"""
        try:
            # 从公平队列获取下一个任务
            # 避免循环导入，直接使用全局实例
            import sys
            if 'core.user_identification' in sys.modules:
                fair_queue = sys.modules['core.user_identification'].fair_queue
                task = fair_queue.get_next_task()
            else:
                task = None
            
            if task:
                # 创建Task对象
                task_obj = Task(
                    task_id=task['task_id'],
                    user_id=task['user_id'],
                    client_ip=task['client_ip'],
                    task_data=task['task_data'],
                    submit_time=task['submit_time']
                )
                
                # 分配任务到服务器
                if self.load_balancer.assign_task(task_obj):
                    logger.info(f"任务已分配: {task['task_id']}")
                    # 更新公平队列中的任务状态
                    if 'core.user_identification' in sys.modules:
                        fair_queue = sys.modules['core.user_identification'].fair_queue
                        fair_queue.update_task_status(
                            task['task_id'], 
                            'processing',
                            start_time=task_obj.start_time,
                            assigned_server=task_obj.assigned_server
                        )
                        # 从队列中移除已分配的任务
                        fair_queue.remove_task(task['task_id'])
                else:
                    logger.warning(f"任务分配失败: {task['task_id']}")
                    # 将任务重新放回队列
                    if 'core.user_identification' in sys.modules:
                        fair_queue = sys.modules['core.user_identification'].fair_queue
                        fair_queue.add_task(
                            task['task_id'],
                            task['client_ip'],
                            task['task_data'],
                            task['submit_time']
                        )
        except Exception as e:
            logger.error(f"处理待处理任务时出错: {e}")
    
    def distribute_task(self, task: Task) -> bool:
        """分发单个任务"""
        return self.load_balancer.assign_task(task)


# 全局实例
load_balancer = LoadBalancer()
task_distributor = TaskDistributor(load_balancer)


if __name__ == "__main__":
    # 测试代码
    logging.basicConfig(level=logging.INFO)
    
    # 创建测试服务器
    server1 = Server("server-001", "ap-beijing", "basic", endpoint="http://server1.example.com")
    server2 = Server("server-002", "ap-nanjing", "standard", endpoint="http://server2.example.com")
    
    load_balancer.add_server(server1)
    load_balancer.add_server(server2)
    
    # 创建测试任务
    task1 = Task("task-001", "user_192.168.1.100_1", "192.168.1.100", {"image": "test.jpg"})
    task2 = Task("task-002", "user_192.168.1.101_1", "192.168.1.101", {"image": "test2.jpg"})
    
    # 分配任务
    load_balancer.assign_task(task1)
    load_balancer.assign_task(task2)
    
    # 查看系统指标
    metrics = load_balancer.get_system_metrics()
    print(f"系统指标: {json.dumps(metrics, indent=2, ensure_ascii=False)}")
    
    # 完成任务
    load_balancer.complete_task("task-001", success=True)
    
    # 查看更新后的指标
    metrics = load_balancer.get_system_metrics()
    print(f"更新后指标: {json.dumps(metrics, indent=2, ensure_ascii=False)}")
